# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import numpy

from experimental.prophet.graph import graph
from experimental.prophet.graph.util import to_data_type_proto
from experimental.prophet.ops.timer import custom_timer


def _sanitize_value(x):
  if isinstance(x, numpy.int64):
    return int(x)
  else:
    return x


def dataframe_feeder(dataframe,
                     timestamp_column='timestamp',
                     timestamp=None,
                     add_timer=True,
                     node_name=None,
                     namespace=None):
  assert timestamp_column in dataframe.columns
  if timestamp is None:
    timestamp = graph.default_timestamp()

  columns = []
  for colname in dataframe.columns:
    dtype = to_data_type_proto(dataframe[colname].dtype)
    columns.append({'name': colname, 'dtype': dtype})

  data = []
  for row_idx in range(dataframe.shape[0]):
    row = [
        _sanitize_value(dataframe.iloc[row_idx, col_idx]) for col_idx in range(dataframe.shape[1])
    ]
    data.append(row)

  attrs = {'columns': columns, 'timestamp_column': timestamp_column, 'data': data}

  node = graph.create_node('DataframeFeeder', [timestamp],
                           attrs=attrs,
                           name=node_name,
                           namespace=namespace)

  if add_timer:
    custom_timer(list(dataframe[timestamp_column]))

  return node.outputs
